﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using System.IO;
using System.Drawing;
using System.Drawing.Drawing2D;
using System.Drawing.Imaging;
using Microsoft.ML.OnnxRuntime;
using Microsoft.ML.OnnxRuntime.Tensors;

namespace OPCService.Service
{
    // 检测结果数据结构
    public class DetectionResult
    {
        public RectangleF Bounds { get; set; } // 边界框
        public float Confidence { get; set; }   // 置信度
        public int ClassId { get; set; }       // 类别ID
        public string Label { get; set; }      // 类别标签
    }

    // 转换参数结构体
    public struct TransformParams
    {
        public float Scale { get; set; }
        public int OffsetX { get; set; }
        public int OffsetY { get; set; }
        public int ScaledWidth { get; set; }
        public int ScaledHeight { get; set; }
    }

    class OnnxDetector : IDisposable
    {
        // 模型配置参数
        private const int InputSize = 640; // YOLOv8 标准输入尺寸
        private const float ConfidenceThreshold = 0.5f; // 置信度阈值
        private const float IoUThreshold = 0.5f; // IoU阈值(NMS使用)

        // ONNX Runtime 推理会话
        private readonly InferenceSession _session;

        // 类别标签
        private string[] _classLabels;

        // 初始化模型
        public OnnxDetector(string modelPath, string[] _classes)
        {
            _classLabels = _classes;
            // 配置Session选项 (可选GPU加速)
            var sessionOptions = new SessionOptions();

            try
            {
                // 启用CUDA加速 (需要安装Microsoft.ML.OnnxRuntime.Gpu)
                sessionOptions.AppendExecutionProvider_CUDA();
                Console.WriteLine("CUDA 执行提供者已启用。");
            }
            catch (Exception ex)
            {
                Console.WriteLine($"CUDA 不可用，原因：{ex.Message}。回退到 CPU。");
                // 如果 CUDA 不可用，不添加任何提供者，ONNX 会默认使用 CPU
            }

            // 启用TensorRT加速
            // sessionOptions.AppendExecutionProvider_Tensorrt();

            _session = new InferenceSession(modelPath, sessionOptions);
        }

        // 执行目标检测
        public List<DetectionResult> Detect(Image image)
        {
            // 获取原始图像尺寸
            int originalWidth = image.Width;
            int originalHeight = image.Height;

            // 预处理图像并获取转换参数
            var (inputTensor, transformParams) = PreprocessImage(image);

            // 准备模型输入
            var inputs = new List<NamedOnnxValue> { NamedOnnxValue.CreateFromTensor("images", inputTensor) };

            // 执行推理
            //using IDisposableReadOnlyCollection<DisposableNamedOnnxValue> results = _session.Run(inputs);
            IDisposableReadOnlyCollection<DisposableNamedOnnxValue> results = _session.Run(inputs);


            // 获取输出数据
            var outputTensor = results.FirstOrDefault()?.AsTensor<float>();
            if (outputTensor == null) return new List<DetectionResult>();

            // 解析模型输出
            var detections = ParseModelOutput(outputTensor, transformParams, originalWidth, originalHeight);

            // 应用非极大值抑制
            return ApplyNMS(detections);
        }

        // 图像预处理和转换参数计算
        private (DenseTensor<float>, TransformParams) PreprocessImage(Image image)
        {
            // 计算缩放比例并保持宽高比
            float scale = Math.Min((float)InputSize / image.Width, (float)InputSize / image.Height);
            int scaledWidth = (int)(image.Width * scale);
            int scaledHeight = (int)(image.Height * scale);

            // 计算填充偏移量
            int offsetX = (InputSize - scaledWidth) / 2;
            int offsetY = (InputSize - scaledHeight) / 2;

            // 创建目标图像(640x640)并填充灰色边框
            //using var scaledImage = new Bitmap(InputSize, InputSize);
            var scaledImage = new Bitmap(InputSize, InputSize);
            using (var g = Graphics.FromImage(scaledImage))
            {
                g.Clear(Color.Gray);
                g.InterpolationMode = InterpolationMode.HighQualityBicubic;
                g.DrawImage(image, offsetX, offsetY, scaledWidth, scaledHeight);
            }

            // 将图像转换为RGB格式并归一化
            var inputTensor = new DenseTensor<float>(new[] { 1, 3, InputSize, InputSize });
            var bitmapData = scaledImage.LockBits(
                new Rectangle(0, 0, InputSize, InputSize),
                ImageLockMode.ReadOnly,
                PixelFormat.Format24bppRgb);

            try
            {
                int bytesPerPixel = 3;
                int stride = bitmapData.Stride;
                byte[] pixelBuffer = new byte[stride * InputSize];
                System.Runtime.InteropServices.Marshal.Copy(bitmapData.Scan0, pixelBuffer, 0, pixelBuffer.Length);

                for (int y = 0; y < InputSize; y++)
                {
                    int rowOffset = y * stride;
                    for (int x = 0; x < InputSize; x++)
                    {
                        int pixelOffset = rowOffset + x * bytesPerPixel;
                        float r = pixelBuffer[pixelOffset + 2] / 255.0f;
                        float g = pixelBuffer[pixelOffset + 1] / 255.0f;
                        float b = pixelBuffer[pixelOffset] / 255.0f;

                        inputTensor[0, 0, y, x] = r;
                        inputTensor[0, 1, y, x] = g;
                        inputTensor[0, 2, y, x] = b;
                    }
                }
            }
            finally
            {
                scaledImage.UnlockBits(bitmapData);
            }

            // 返回处理后的张量和转换参数
            var transformParams = new TransformParams
            {
                Scale = scale,
                OffsetX = offsetX,
                OffsetY = offsetY,
                ScaledWidth = scaledWidth,
                ScaledHeight = scaledHeight
            };

            return (inputTensor, transformParams);
        }

        // 解析模型输出 - 修复Y坐标偏移问题
        private List<DetectionResult> ParseModelOutput(Tensor<float> output, TransformParams transformParams, int originalWidth, int originalHeight)
        {
            var detections = new List<DetectionResult>();
            int numClasses = _classLabels.Length;
            int numBoxes = output.Dimensions[2];

            // 获取转换参数
            float scale = transformParams.Scale;
            int offsetX = transformParams.OffsetX;
            int offsetY = transformParams.OffsetY;
            int scaledWidth = transformParams.ScaledWidth;
            int scaledHeight = transformParams.ScaledHeight;

            for (int i = 0; i < numBoxes; i++)
            {
                // 获取边界框数据
                float xCenter = output[0, 0, i];
                float yCenter = output[0, 1, i];
                float width = output[0, 2, i];
                float height = output[0, 3, i];

                // 获取最大置信度的类别
                float maxConfidence = 0;
                int classId = -1;

                for (int c = 0; c < numClasses; c++)
                {
                    float confidence = output[0, 4 + c, i];
                    if (confidence > maxConfidence)
                    {
                        maxConfidence = confidence;
                        classId = c;
                    }
                }

                if (maxConfidence < ConfidenceThreshold) continue;

                // 将中心坐标转换为角坐标
                float x1 = xCenter - width / 2;
                float y1 = yCenter - height / 2;
                float x2 = xCenter + width / 2;
                float y2 = yCenter + height / 2;

                // 关键修复：将坐标从640x640空间转换回原始图像空间
                // 1. 减去填充偏移量
                x1 = (x1 - offsetX) / scale;
                y1 = (y1 - offsetY) / scale;
                x2 = (x2 - offsetX) / scale;
                y2 = (y2 - offsetY) / scale;

                // 2. 限制在原始图像范围内
                x1 = Math.Max(0, Math.Min(x1, originalWidth));
                y1 = Math.Max(0, Math.Min(y1, originalHeight));
                x2 = Math.Max(0, Math.Min(x2, originalWidth));
                y2 = Math.Max(0, Math.Min(y2, originalHeight));

                // 创建检测结果
                detections.Add(new DetectionResult
                {
                    Bounds = new RectangleF(
                        x1, y1,
                        x2 - x1,
                        y2 - y1),
                    Confidence = maxConfidence,
                    ClassId = classId,
                    Label = _classLabels[classId]
                });
            }

            return detections;
        }

        // 非极大值抑制(NMS)算法
        private List<DetectionResult> ApplyNMS(List<DetectionResult> detections)
        {
            detections = detections.OrderByDescending(d => d.Confidence).ToList();
            var finalDetections = new List<DetectionResult>();

            while (detections.Count > 0)
            {
                var current = detections[0];
                finalDetections.Add(current);
                detections.RemoveAt(0);

                for (int i = detections.Count - 1; i >= 0; i--)
                {
                    float iou = CalculateIoU(current.Bounds, detections[i].Bounds);
                    if (iou > IoUThreshold)
                    {
                        detections.RemoveAt(i);
                    }
                }
            }

            return finalDetections;
        }

        // 计算交并比
        private float CalculateIoU(RectangleF rect1, RectangleF rect2)
        {
            float x1 = Math.Max(rect1.Left, rect2.Left);
            float y1 = Math.Max(rect1.Top, rect2.Top);
            float x2 = Math.Min(rect1.Right, rect2.Right);
            float y2 = Math.Min(rect1.Bottom, rect2.Bottom);

            float intersection = Math.Max(0, x2 - x1) * Math.Max(0, y2 - y1);
            float area1 = rect1.Width * rect1.Height;
            float area2 = rect2.Width * rect2.Height;

            return intersection / (area1 + area2 - intersection);
        }

        // 在图像上绘制检测结果
        public Bitmap DrawResults(Image originalImage, List<DetectionResult> results)
        {
            var bitmap = new Bitmap(originalImage);
            using (var g = Graphics.FromImage(bitmap))
            {
                var font = new Font("Arial", 24, FontStyle.Bold);
                var brush = new SolidBrush(Color.White);
                var penWidth = Math.Max(originalImage.Width / 400, 2);

                foreach (var detection in results)
                {
                    var color = GetClassColor(detection.ClassId);
                    //using var pen = new Pen(color, penWidth);
                    var pen = new Pen(color, penWidth);

                    // 绘制边界框
                    var rect = detection.Bounds;
                    g.DrawRectangle(pen, rect.X, rect.Y, rect.Width, rect.Height);

                    // 绘制标签背景
                    string label = $"{detection.Label}: {detection.Confidence:F2}";
                    var labelSize = g.MeasureString(label, font);
                    float labelY = Math.Max(0, rect.Y - labelSize.Height);

                    g.FillRectangle(new SolidBrush(Color.FromArgb(200, color)),
                        rect.X, labelY,
                        labelSize.Width, labelSize.Height);

                    // 绘制标签文本
                    g.DrawString(label, font, brush, rect.X, labelY);
                }
            }

            return bitmap;
        }

        // 为每个类别生成颜色
        private Color GetClassColor(int classId)
        {
            // 使用HSV颜色空间确保颜色区分度高
            float hue = (classId * 137.5f) % 360; // 黄金角间隔
            return HsvToRgb(hue, 0.9f, 0.9f);
        }

        // HSV转RGB
        private Color HsvToRgb(float hue, float saturation, float value)
        {
            int hi = (int)(hue / 60) % 6;
            float f = hue / 60 - (int)(hue / 60);

            float v = value * 255;
            float p = v * (1 - saturation);
            float q = v * (1 - f * saturation);
            float t = v * (1 - (1 - f) * saturation);

            float r, g, b;
            switch (hi)
            {
                case 0:
                    r = v;
                    g = t;
                    b = p;
                    break;
                case 1:
                    r = q;
                    g = v;
                    b = p;
                    break;
                case 2:
                    r = p;
                    g = v;
                    b = t;
                    break;
                case 3:
                    r = p;
                    g = q;
                    b = v;
                    break;
                case 4:
                    r = t;
                    g = p;
                    b = v;
                    break;
                default: // case 5
                    r = v;
                    g = p;
                    b = q;
                    break;
            }

            return Color.FromArgb((int)r, (int)g, (int)b);
        }

        // 释放资源
        public void Dispose()
        {
            _session?.Dispose();
        }
    }
}
